use std::cmp::Ordering;

#[derive(Debug)]
struct AVLTreeNode<T> {
    key: T,
    value: T,
    height: usize,
    left: Option<Box<AVLTreeNode<T>>>,
    right: Option<Box<AVLTreeNode<T>>>,
}

impl<T: Ord + Copy> AVLTreeNode<T> {
    fn new(key: T, value: T) -> Self {
        AVLTreeNode {
            key,
            value,
            height: 1,
            left: None,
            right: None,
        }
    }

    fn height(&self) -> usize {
        self.height
    }

    fn balance_factor(&self) -> isize {
        let left_height = self.left.as_ref().map_or(0, |node| node.height());
        let right_height = self.right.as_ref().map_or(0, |node| node.height());
        (left_height as isize) - (right_height as isize)
    }

    fn rotate_right(mut root: Box<AVLTreeNode<T>>) -> Box<AVLTreeNode<T>> {
        let mut left_child = root.left.take().unwrap();
        root.left = left_child.right.take();
        left_child.right = Some(root);
        left_child
    }

    fn rotate_left(mut root: Box<AVLTreeNode<T>>) -> Box<AVLTreeNode<T>> {
        let mut right_child = root.right.take().unwrap();
        root.right = right_child.left.take();
        right_child.left = Some(root);
        right_child
    }

    fn balance(mut self) -> Box<AVLTreeNode<T>> {
        match self.balance_factor() {
            -2 => {
                if self.right.as_ref().unwrap().balance_factor() == 1 {
                    let right_child = self.right.take().unwrap();
                    self.right = Some(Self::rotate_right(right_child));
                }
                Self::rotate_left(Box::new(self))
            }
            2 => {
                if self.left.as_ref().unwrap().balance_factor() == -1 {
                    let left_child = self.left.take().unwrap();
                    self.left = Some(Self::rotate_left(left_child));
                }
                Self::rotate_right(Box::new(self))
            }
            _ => Box::new(self),
        }
    }

    fn insert(mut self, key: T, value: T) -> Box<AVLTreeNode<T>> {
        match key.cmp(&self.key) {
            Ordering::Less => {
                if let Some(left) = self.left.take() {
                    self.left = Some(left.insert(key, value));
                } else {
                    self.left = Some(Box::new(AVLTreeNode::new(key, value)));
                }
            }
            Ordering::Greater => {
                if let Some(right) = self.right.take() {
                    self.right = Some(right.insert(key, value));
                } else {
                    self.right = Some(Box::new(AVLTreeNode::new(key, value)));
                }
            }
            Ordering::Equal => {
                self.value = value;
            }
        }

        self.balance()
    }

    fn lookup(&self, key: T) -> Option<T> {
        match key.cmp(&self.key) {
            Ordering::Less => self.left.as_ref().and_then(|node| node.lookup(key)),
            Ordering::Greater => self.right.as_ref().and_then(|node| node.lookup(key)),
            Ordering::Equal => Some(self.value),
        }
    }

    fn to_array(&self, array: &mut Vec<T>) {
        if let Some(ref left) = self.left {
            left.to_array(array);
        }
        array.push(self.key);
        if let Some(ref right) = self.right {
            right.to_array(array);
        }
    }
}

#[derive(Debug)]
pub struct AVLTree<T> {
    root: Option<Box<AVLTreeNode<T>>>,
}

impl<T: Ord + Copy> AVLTree<T> {
    pub fn new() -> Self {
        AVLTree { root: None }
    }

    pub fn insert(&mut self, key: T, value: T) {
        if let Some(ref mut root) = self.root {
            *root = root.insert(key, value);
        } else {
            self.root = Some(Box::new(AVLTreeNode::new(key, value)));
        }
    }

    pub fn lookup(&self, key: T) -> Option<T> {
        self.root.as_ref().and_then(|node| node.lookup(key))
    }

    pub fn to_array(&self) -> Vec<T> {
        let mut array = Vec::new();
        if let Some(ref root) = self.root {
            root.to_array(&mut array);
        }
        array
    }
}

fn main() {
    let mut tree = AVLTree::new();
    tree.insert(10, 10);
    tree.insert(20, 20);
    tree.insert(30, 30);
    tree.insert(40, 40);
    tree.insert(50, 50);

    println!("{:?}", tree.lookup(30));
    println!("{:?}", tree.to_array());
}
